import boto3
import psycopg2
from wt.connector.db_manager.db_untils_tools import DBUntilsTools
from wt.connector.data_connect import DataConnect

class DBS3RedshiftManager(DataConnect, DBUntilsTools):

    def __init__(self):
        super().__init__()

    def create_ddl(self):
        col = {
            "id": "int4 NOT NULL,",
            "c_02_integer_col": "int4,",
            "c_03_double_col": "float8,",
            "c_04_float_col": "float4,",
            "c_05_smallint_col": "int2,",
            "c_06_bigint_col": "int8,",
            "c_07_tinyint_col": "int2,",
            "c_08_decimal_col": "numeric(10,2),",
            "c_09_varchar_col": "varchar(50),",
            "c_10_date_col": "date,",
            "c_11_datetime_col": "timestamp,",
            "PRIMARY KEY": "(id)"
        }
        return col

    def generate_create_ddl(self, table):
        col = self.create_ddl()
        ddl_sql = f'CREATE TABLE {table} ( '
        for key in col:
            line = key + " " + col[key] + "\n"
            ddl_sql += line
        ddl_sql += ');'
        return ddl_sql

    def general_col_value(self, table):
        c_02_integer_val = self.get_random(-2147483648, 2147483648)
        c_03_double_val = self.get_random_float(-500, 500)
        c_04_float_val = self.get_random_float(-500, 500)
        c_05_smallint_val = self.get_random(-32768, 32767)
        c_06_bigint_val = self.get_random(-9223372036854775808, 9223372036854775807)
        c_07_tinyint_val = self.get_random(-128, 123)
        c_08_decimal_val = self.get_random_float(-500, 500)
        c_09_varchar_val = self.get_random_string(8)
        c_10_date_val = self.get_random_date()
        c_11_datetime_val = self.get_random_datetime()
        sql = f"INSERT INTO {table} (c_02_integer_col, c_03_double_col, c_04_float_col, c_05_smallint_col, c_06_bigint_col," \
              f" c_07_tinyint_col, c_08_decimal_col, c_09_varchar_col, c_10_date_col, c_11_datetime_col) " \
              f"VALUES ({c_02_integer_val},{c_03_double_val},{c_04_float_val},{c_05_smallint_val},{c_06_bigint_val}," \
              f"{c_07_tinyint_val},{c_08_decimal_val},'{c_09_varchar_val}','{c_10_date_val}','{c_11_datetime_val}');"
        return sql

    def generate_dml_sql(self, table, lines):
        insert_sql = ""
        while lines > 0:
            tmp = self.general_col_value(table)
            insert_sql += tmp + "\n"
            lines -= 1
        return insert_sql

    def execute_create_and_insert_sql(self, db, table, num, is_recreate=None):
        if is_recreate:
            ddl_sql = self.generate_create_ddl(table)
            self.exec_s3redshift_sql(db, f"DROP TABLE IF EXISTS {table};")
            self.exec_s3redshift_sql(db, ddl_sql)
        insert_sql = self.generate_dml_sql(table, num)
        for line in insert_sql.strip().split("\n"):
            self.exec_s3redshift_sql(db, line)

    def exec_customer_sql(self, db, sql):
        res = self.exec_s3redshift_sql(db, sql)
        return res

    def load_data_from_s3(self, table, s3_path, iam_role):
        copy_sql = f"""
        COPY {table}
        FROM '{s3_path}'
        IAM_ROLE '{iam_role}'
        FORMAT AS CSV;
        """
        self.exec_s3redshift_sql(sql_str=copy_sql)

    def download_latest_file(self, tableName, local_path):
        prefix = f"{tableName}/"
        try:
            response = self.s3_client.list_objects_v2(Bucket=self.s3_bucket, Prefix=prefix)
            if 'Contents' not in response:
                print("No files found.")
                return False

            files = response['Contents']
            latest_file = max(files, key=lambda x: x['LastModified'])
            latest_file_key = latest_file['Key']

            self.s3_client.download_file(self.s3_bucket, latest_file_key, local_path)
            print(f"Downloaded {latest_file_key} to {local_path}")
            return True
        except Exception as e:
            print(f"An error occurred: {e}")
            return False

if __name__ == '__main__':
    manager = DBS3RedshiftManager()
    manager.execute_create_and_insert_sql("dev", "auto_test_20240527", 10, is_recreate=True)
